#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created by PyCharm.

@Date    : Thu Feb 25 2021 
@Time    : 04:48:49
@File    : model.py
@Author  : alpha
"""


import torch
from torch import nn

from src.regvgg import RepVGG18
from src.resnet import ConvBN, ResNet18



class MLP(nn.Module):

    def __init__(self, in_channels):
        super(MLP, self).__init__()
        self.c1 = ConvBN(in_channels, 256, kernel_size=1, padding=0)
        self.c2 = ConvBN(256, 256, kernel_size=1, padding=0)
        self.cls = nn.Conv2d(256, 1, kernel_size=1)

    def forward(self, feat):
        c1 = self.c1(feat)
        c2 = self.c2(c1)
        cls = self.cls(c2)
        return cls


class FaceSpoofNet(nn.Module):

    def __init__(self, feat_channels, backbone):
        super(FaceSpoofNet, self).__init__()
        self.feat_net = backbone(out_channels=feat_channels)
        self.cls_net = MLP(in_channels=feat_channels)

    def feat_forward(self, x, y, z):
        x = self.feat_net(x)
        y = self.feat_net(y)
        z = self.feat_net(z)
        return x, y, z

    def cls_forward(self, x, y, z):
        x = self.cls_net(x)
        y = self.cls_net(y)
        z = self.cls_net(z)
        return x, y, z

    def forward(self, x):
        # ONLY FOR EVALUATION
        x = self.feat_net(x)
        x = self.cls_net(x)
        return torch.sigmoid(x)